import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, num_heads, num_groups):
        super(GroupedQueryAttention, self).__init__()
        assert d_model % num_heads == 0
        assert num_heads % num_groups == 0
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.num_groups = num_groups
        self.group_dim = self.head_dim * num_groups
        self.d_model = d_model        
        # Linear layers for Query, Key, and Value
        self.wq = nn.Linear(d_model, d_model)
        self.wk = nn.Linear(d_model, self.group_dim)  # Shared Key per group
        self.wv = nn.Linear(d_model, self.group_dim)  # Shared Value per group        
        self.fc = nn.Linear(d_model, d_model)   
    def split_heads(self, x, batch_size):
        x = x.view(batch_size, -1, self.num_heads, self.head_dim)
        return x.permute(0, 2, 1, 3)  # (batch_size, num_heads, seq_len, head_dim)    
    def forward(self, x, mask=None):
        batch_size = x.size(0)      
        # Linear projections
        q = self.wq(x)  # (batch_size, seq_len, d_model)
        k = self.wk(x)  # (batch_size, seq_len, group_dim)
        v = self.wv(x)  # (batch_size, seq_len, group_dim)      
        # Split Query into heads
        q = self.split_heads(q, batch_size)  # (batch_size, num_heads, seq_len, head_dim)        
        # Split Key and Value into groups and expand within the group
        k = k.view(batch_size, -1, self.num_groups, self.head_dim)
        v = v.view(batch_size, -1, self.num_groups, self.head_dim)        
        k = k.repeat(1, 1, 1, self.num_heads // self.num_groups)
        v = v.repeat(1, 1, 1, self.num_heads // self.num_groups)      
        k = k.view(batch_size, -1, self.num_heads, self.head_dim)
        v = v.view(batch_size, -1, self.num_heads, self.head_dim)        
        # Scaled Dot-Product Attention
        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim)  # (batch_size, num_heads, seq_len, seq_len)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)        
        attention_weights = F.softmax(scores, dim=-1)  # (batch_size, num_heads, seq_len, seq_len)
        output = torch.matmul(attention_weights, v)  # (batch_size, num_heads, seq_len, head_dim)      
        # Concatenate heads
        output = output.permute(0, 2, 1, 3).contiguous()  # (batch_size, seq_len, num_heads, head_dim)
        output = output.view(batch_size, -1, self.d_model)  # (batch_size, seq_len, d_model)       
        # Final linear layer
        output = self.fc(output)  # (batch_size, seq_len, d_model        
        return output, attention_weights

    def incremental_forward_kvcache(self, x_t, past_k=None, past_v=None):
        """
        增量推理（带KV缓存）。
        输入: x_t (B, T_new, d_model)；
        缓存: past_k/past_v 为 (B, G, T_past, head_dim)。
        返回: output_t (B, T_new, d_model), attn (B, H, T_new, T_total), new_past_k, new_past_v。
        """
        batch_size = x_t.size(0)
        device = x_t.device

        # Projections
        q = self.wq(x_t)
        k = self.wk(x_t)
        v = self.wv(x_t)

        # Split Query into heads -> (B, H, T_new, D)
        q = self.split_heads(q, batch_size)

        # Split K/V into groups -> (B, G, T_new, D)
        k = k.view(batch_size, -1, self.num_groups, self.head_dim).permute(0, 2, 1, 3).contiguous()
        v = v.view(batch_size, -1, self.num_groups, self.head_dim).permute(0, 2, 1, 3).contiguous()

        # Concat cache along time
        if past_k is not None and past_v is not None:
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)

        # Map heads to groups to build per-head K/V (B, H, T_total, D)
        heads_per_group = self.num_heads // self.num_groups
        head_to_group = (torch.arange(self.num_heads, device=device) // heads_per_group)
        k_h = k[:, head_to_group, :, :]
        v_h = v[:, head_to_group, :, :]

        # Attention over cached keys
        scores = torch.matmul(q, k_h.transpose(-1, -2)) / math.sqrt(self.head_dim)
        attn = F.softmax(scores, dim=-1)
        out_h = torch.matmul(attn, v_h)

        # Merge heads
        out = out_h.permute(0, 2, 1, 3).contiguous().view(batch_size, -1, self.d_model)
        out = self.fc(out)

        return out, attn, k, v
